__all__ = ['initialise_cuda', 'MAXGPU', 'CANUSEGPU', 'set_gpu_device', 'close_cuda']
    
try:
    from debugtools import *
    import atexit
    import pycuda
    import pycuda.driver as drv
    pycuda.context = None
    pycuda.isinitialised = False
    MAXGPU = 0 # drv.Device.count()
    CANUSEGPU = True
    
    def initialise_cuda():
        global MAXGPU
        if not pycuda.isinitialised:
            drv.init()
            pycuda.isinitialised = True
            MAXGPU = drv.Device.count()
            log.debug("Initializing PyCUDA, %d GPU(s) found" % MAXGPU)
            return MAXGPU
    
    def set_gpu_device(n):
        """
        This function makes pycuda use GPU number n in the system.
        """
        initialise_cuda()
        log.debug("Setting PyCUDA context number %d" % n)
        try:
            pycuda.context.detach()
        except:
            pass
        if n<MAXGPU:
            pycuda.context = drv.Device(n).make_context()
        else:
            pycuda.context = drv.Device(MAXGPU-1).make_context()
            log.warn("Unable to set GPU device %d, setting device %d instead" % (n, MAXGPU-1))
    
    def close_cuda():
        """
        Closes the current context. MUST be called at the end of the script.
        """
        if pycuda.context is not None:
            try:
                pycuda.context.pop()
                log.debug("Closing current PyCUDA context")
                pycuda.context = None
            except:
                pass
    
    atexit.register(close_cuda)

except:
    MAXGPU = 0
    CANUSEGPU = False

    def initialise_cuda():
        pass
    
    def set_gpu_device(n):
        log.warn("PyCUDA not available")
        pass

    def close_cuda():
        pass

